import 'package:camera/camera.dart';
import 'package:tflite_v2/tflite_v2.dart';
import 'package:logging/logging.dart';
import '../../config/app_config.dart';
import 'dart:typed_data';
import 'package:image/image.dart' as img;

/// TFLite服务
/// 负责TFLite模型的加载和推理
/// 创建时间：2024-01-23

class TFLiteService {
  final _logger = Logger('TFLiteService');
  bool _isModelLoaded = false;
  bool _isInferencing = false; // 添加标志以跟踪推理状态

  /// 模型输入尺寸
  final List<int> _inputShape = [1, 224, 224, 3];

  /// 加载TFLite模型
  Future<void> loadModel() async {
    try {
      String? res = await Tflite.loadModel(
        model: AppConfig.modelPath,
        labels: AppConfig.labelsPath,
      );
      _isModelLoaded = true;
      _logger.info("模型加载成功: $res");
    } catch (e) {
      _logger.severe("模型加载失败: $e");
      _isModelLoaded = false;
    }
  }

  /// 执行姿态估计
  /// 参数：
  /// - image: 相机图像
  /// 返回：
  /// - 检测到的关键点列表
  Future<List<dynamic>?> runPoseEstimation(CameraImage image) async {
    if (!_isModelLoaded) {
      _logger.warning("模型未加载");
      return null;
    }

    if (_isInferencing) {
      _logger.info("模型正在推理中，跳过当前帧");
      return null;
    }

    try {
      _isInferencing = true;

      // 预处理图像
      Uint8List imageBytes = _preprocessCameraImage(image);

      // 检查输入大小是否匹配
      if (imageBytes.length !=
          _inputShape[1] * _inputShape[2] * _inputShape[3] * 4) {
        _logger.warning(
            '输入大小不匹配: 期望 ${_inputShape[1] * _inputShape[2] * _inputShape[3] * 4} 字节, 实际 ${imageBytes.length} 字节');
        return null;
      }

      // 执行模型推理
      var rawResults = await Tflite.runModelOnBinary(
        binary: imageBytes.buffer.asUint8List(),
        numResults: AppConfig.numResults,
        threshold: AppConfig.detectionThreshold,
      );

      return rawResults
          ?.map((result) => {
                'x': (result['x'] as num).toDouble(),
                'y': (result['y'] as num).toDouble(),
              })
          .toList();
    } catch (e, stackTrace) {
      _logger.severe('TFLite推理错误', e, stackTrace);
      return null;
    } finally {
      _isInferencing = false;
    }
  }

  /// 预处理相机图像
  /// 将YUV格式转换为RGB格式，并调整尺寸
  Uint8List _preprocessCameraImage(CameraImage image) {
    try {
      // 将YUV格式转换为RGB格式
      img.Image rgbImage = _convertYUV420toRGB(image);

      // 调整图像尺寸以匹配模型输入
      img.Image resizedImage = img.copyResize(
        rgbImage,
        width: _inputShape[1],
        height: _inputShape[2],
        interpolation: img.Interpolation.linear,
      );

      // 转换为归一化的浮点数组
      var convertedBytes =
          Uint8List(_inputShape[1] * _inputShape[2] * _inputShape[3]);
      int pixelIndex = 0;

      for (var y = 0; y < resizedImage.height; y++) {
        for (var x = 0; x < resizedImage.width; x++) {
          var pixel = resizedImage.getPixel(x, y);
          // 直接使用0-255范围的值
          convertedBytes[pixelIndex++] = pixel.r.toInt();
          convertedBytes[pixelIndex++] = pixel.g.toInt();
          convertedBytes[pixelIndex++] = pixel.b.toInt();
        }
      }

      _logger.info('预处理后图像大小: ${convertedBytes.length * 4} bytes');
      return convertedBytes;
    } catch (e, stackTrace) {
      _logger.severe('图像预处理失败', e, stackTrace);
      rethrow;
    }
  }

  /// 将YUV420格式转换为RGB格式
  img.Image _convertYUV420toRGB(CameraImage image) {
    final int width = image.width;
    final int height = image.height;
    final img.Image rgbImage = img.Image(width: width, height: height);

    final yPlane = image.planes[0];
    final uPlane = image.planes[1];
    final vPlane = image.planes[2];

    final int uvRowStride = uPlane.bytesPerRow;
    final int uvPixelStride = uPlane.bytesPerPixel!;

    for (int y = 0; y < height; y++) {
      for (int x = 0; x < width; x++) {
        final int uvIndex = uvPixelStride * (x ~/ 2) + uvRowStride * (y ~/ 2);
        final int index = y * width + x;

        final yp = yPlane.bytes[index];
        final up = uPlane.bytes[uvIndex];
        final vp = vPlane.bytes[uvIndex];

        // YUV到RGB的转换
        int r = (yp + 1.370705 * (vp - 128)).toInt().clamp(0, 255);
        int g = (yp - 0.698001 * (vp - 128) - 0.337633 * (up - 128))
            .toInt()
            .round()
            .clamp(0, 255);
        int b = (yp + 1.732446 * (up - 128)).toInt().clamp(0, 255);

        rgbImage.setPixelRgba(x, y, r, g, b, 255);
      }
    }

    return rgbImage;
  }

  /// 释放资源
  void dispose() {
    if (_isModelLoaded) {
      Tflite.close();
      _isModelLoaded = false;
    }
  }
}
